package com.cy.embeddingdemo;

import com.cy.embeddingdemo.config.AiConfig;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.ClassPathDocumentLoader;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentByCharacterSplitter;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.ollama.OllamaChatModel;
import dev.langchain4j.model.ollama.OllamaEmbeddingModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import jakarta.annotation.Resource;
import org.junit.jupiter.api.Test;
import org.springframework.boot.test.context.SpringBootTest;

import java.util.List;

@SpringBootTest
public class Rag {

    @Resource
    public OpenAiEmbeddingModel embeddingModel;



    @Test
    public void readText() {
        Document document = ClassPathDocumentLoader.loadDocument("rag/xiaoshuo.txt", new TextDocumentParser());

        DocumentByCharacterSplitter documentByCharacterSplitter = new DocumentByCharacterSplitter(
                50,  //每段最大字数
                10 // 自然语言重叠最大字数
        );

        List<TextSegment> segments = documentByCharacterSplitter.split(document);

        List<Embedding> embeddings = embeddingModel.embedAll(segments).content();


        // 存入
        InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();
        embeddingStore.addAll(embeddings,segments);

        // 生成向量
        Response<Embedding> embed = embeddingModel.embed("白发老人在课堂里干嘛");
        EmbeddingSearchRequest build = EmbeddingSearchRequest.builder().queryEmbedding(embed.content()).build();
        // 查询
        EmbeddingSearchResult<TextSegment> results = embeddingStore.search(build);
        for (EmbeddingMatch<TextSegment> match : results.matches()) {
            System.out.println(match.embedded().text() + ",分数为：" + match.score());

        }
    }

    @Resource
    public OllamaChatModel ollamaChatModel;

    @Test
    public void rag() {


        Document document = ClassPathDocumentLoader.loadDocument("rag/xiaoshuo.txt", new TextDocumentParser());

        DocumentByCharacterSplitter documentByCharacterSplitter = new DocumentByCharacterSplitter(
                50,  //每段最大字数
                10 // 自然语言重叠最大字数
        );

        List<TextSegment> segments = documentByCharacterSplitter.split(document);

        List<Embedding> embeddings = embeddingModel.embedAll(segments).content();


        // 存入
        InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();
        embeddingStore.addAll(embeddings,segments);





        ContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(embeddingModel)
                .maxResults(5) // 最相似的5个结果
                .minScore(0.6) // 只找相似度在0.6以上的内容
                .build();

        // 为Assistant动态代理对象  chat  --->  对话内容存储ChatMemory----> 聊天记录ChatMemory取出来 ---->放入到当前对话中
        AiConfig.Assistant assistant = AiServices.builder(AiConfig.Assistant.class)
                .chatLanguageModel(ollamaChatModel)
                .contentRetriever(contentRetriever)
                .build();

        System.out.println(assistant.chat("白发老人的身份"));




    }




}
